import torch
import torch.nn as nn

class MultiScaleCNN(nn.Module):
    def __init__(self, config, out_channels=32):
        super().__init__()
        self.conv_blocks = nn.ModuleList([
            nn.Sequential(
                nn.Conv1d(config.INPUT_DIM, out_channels, kernel_size=3, padding=1),
                nn.BatchNorm1d(out_channels),
                nn.ReLU(),
            ),
            nn.Sequential(
                nn.Conv1d(config.INPUT_DIM, out_channels, kernel_size=5, padding=2),
                nn.BatchNorm1d(out_channels),
                nn.ReLU(),
            ),
            nn.Sequential(
                nn.Conv1d(config.INPUT_DIM, out_channels, kernel_size=7, padding=3),
                nn.BatchNorm1d(out_channels),
                nn.ReLU(),
            )
        ])
    
    def forward(self, x):
        x = x.permute(0, 2, 1)  # [batch, features, seq_len]
        
        cnn_outs = []
        for conv in self.conv_blocks:
            cnn_outs.append(conv(x))
            
        out = torch.cat(cnn_outs, dim=1)  # [batch, out_channels*3, seq_len]
        return out.permute(0, 2, 1)  # [batch, seq_len, out_channels*3] 